Solutions/Vectra XDR/Data Connectors/VectraDataConnector/SharedCode/collector.py

"""This file contains methods for validations, checkpoint, pulling and pushing data."""

import sys
import datetime
import json
import inspect
import requests
import hashlib
from datetime import timedelta
from requests.auth import HTTPBasicAuth
from ..SharedCode import consts, keyvault_secrets_management
from ..SharedCode.azure_sentinel import AzureSentinel


class BaseCollector:
    """This class contains methods to create object and helper methods."""

    def __init__(self, applogger, function_name, client_id, client_secret) -> None:
        """Initialize instance variable for class."""
        self.connection_string = consts.CONNECTION_STRING
        self.base_url = consts.BASE_URL
        self.client_id = client_id
        self.client_secret = client_secret
        self.start_time = consts.START_TIME
        self.applogger = applogger
        self.azuresentinel = AzureSentinel()
        self.session = requests.Session()
        self.session.headers["User-Agent"] = consts.USER_AGENT
        self.session.headers["Content-Type"] = "application/x-www-form-urlencoded"
        self.function_name = function_name
        self.access_token_expiration = None
        self.refresh_token_expiration = None
        self.access_token = None
        self.refresh_token = None
        self.keyvault_obj = keyvault_secrets_management.KeyVaultSecretManage(applogger)

    def load_tokens_from_keyvault(self):
        """Load tokens from KeyVault based on access and refresh token keys."""
        properties_list = self.keyvault_obj.get_properties_list_of_secrets()
        if self.access_token_key in properties_list:
            self.access_token = self.keyvault_obj.get_keyvault_secret(self.access_token_key)
        if self.refresh_token_key in properties_list:
            self.refresh_token = self.keyvault_obj.get_keyvault_secret(self.refresh_token_key)

    def load_expiration_time_from_checkpoint_file(self):
        """Load expiration time from checkpoint file."""
        __method_name = inspect.currentframe().f_code.co_name
        try:
            data = self.state.get()
            if data is not None and data:
                data = json.loads(data)
                self.access_token_expiration = data.get(self.access_token_expiry, None)
                self.refresh_token_expiration = data.get(self.refresh_token_expiry, None)
        except Exception as err:
            self.applogger.error(
                "{}(method={}) : {} : Exception occured while getting data from storage account. : {}".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, err
                )
            )
            raise Exception(f"Exception occured during getting data from azure : {err}.")

    def add_token_expiration_to_checkpoint_file(self, token_data, from_refresh_token: bool):
        """Add token to keyvault and expiration time to checkpoint file.

        Args:
            token_data (json): token_data to be added in checkpoint file and keyvault.
        """
        try:
            __method_name = inspect.currentframe().f_code.co_name
            checkpoint_data = {}
            self.access_token = token_data["access_token"]
            self.session.headers["Authorization"] = "bearer {}".format(
                    self.access_token
                )
            self.access_token_expiration = (
                datetime.datetime.now() + timedelta(seconds=token_data["expires_in"])
            ).isoformat()
            if not from_refresh_token:
                self.applogger.info(
                    "{}(method={}) : {} : Both access token and refresh token are generated successfully.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
                self.refresh_token = token_data["refresh_token"]
                self.keyvault_obj.set_keyvault_secret(self.refresh_token_key, self.refresh_token)
                self.refresh_token_expiration = (
                    datetime.datetime.now() + timedelta(seconds=token_data["refresh_expires_in"])
                ).isoformat()
                token_data = {
                    self.access_token_expiry: self.access_token_expiration,
                    self.refresh_token_expiry: self.refresh_token_expiration,
                }
            else:
                self.applogger.info(
                    "{}(method={}) : {} :  Access Token Generated from Refresh Token.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
                token_data = {self.access_token_expiry: self.access_token_expiration}
            checkpoint_data.update(token_data)
            self.keyvault_obj.set_keyvault_secret(self.access_token_key, self.access_token)
            self.save_checkpoint(data=checkpoint_data, token_expiry_time=True)
        except Exception as err:
            self.applogger.error(
                "{}(method={}) : {} : Error occured in storing token or expiration time: {} ".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, err
                )
            )
            raise Exception(
                "Error occured in storing token to keyvault or expiration time to checkpoint file : {}".format(err)
            )

    def generate_access_token(self, from_refresh_token: bool):
        """Generate the Token and Refresh token.

        Args:
            from_refresh_token (bool): To generate Access Token from refresh token or not

        Raises:
            ValueError: When response status code is not [200,201]
        """
        try:
            __method_name = inspect.currentframe().f_code.co_name
            url = self.base_url + consts.OAUTH2_ENDPOINT
            auth = HTTPBasicAuth(self.client_id, self.client_secret)
            data = None
            if from_refresh_token:
                data = {
                    "grant_type": "refresh_token",
                    "refresh_token": self.refresh_token,
                }
                self.applogger.debug(
                    "{}(method={}) : {} : Generating access token using refresh token.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
            else:
                data = {"grant_type": "client_credentials"}
                self.applogger.debug(
                    "{}(method={}) : {} : Generating access token.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
            response = requests.post(url, auth=auth, data=data, timeout=consts.API_TIMEOUT)

            if response.status_code in [200, 201]:
                token_data = response.json()
                self.add_token_expiration_to_checkpoint_file(token_data, from_refresh_token)
                self.update_checkpoint_of_disabling_function(make_count_zero=True)
                self.applogger.debug(
                    "{}(method={}) : {} : New access token generated successfully ".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
            elif response.status_code in [400, 401, 403, 404]:
                self.update_checkpoint_of_disabling_function()
                self.applogger.error(
                    "{}(method={}) : {} : status code={}, reason={}".format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        response.status_code,
                        response.text
                    )
                )
                raise Exception
            else:
                self.applogger.error(
                    "{}(method={}) : {} : Unknown status code: {} ".format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        response.status_code,
                    )
                )
                raise Exception

        except requests.exceptions.RequestException as err:
            self.applogger.error(
                "{}(method={}) : {} : Error occurred while generating a Access Token : {} ".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, err
                )
            )
            raise Exception(f"Error occurred while generating a Access Token : {err}")
        except Exception as err:
            self.applogger.error(
                "{}(method={}) : {} : Error occurred while generating a Access Token : {} ".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, err
                )
            )
            raise Exception(err)

    def validate_access_token(self):
        """Validate the existing token and refresh token."""
        try:
            __method_name = inspect.currentframe().f_code.co_name
            from_refresh_token = 0
            if self.access_token:
                if self.is_token_expired(self.access_token_expiration) and not self.is_token_expired(
                    self.refresh_token_expiration
                ):
                    self.applogger.info(
                        "{}(method={}) : {} : Access token expired. Generating new access token using refresh token.".format(
                            consts.LOGS_STARTS_WITH, __method_name, self.function_name
                        )
                    )
                    from_refresh_token = 1
                    self.generate_access_token(from_refresh_token)
                elif self.is_token_expired(self.access_token_expiration) and self.is_token_expired(
                    self.refresh_token_expiration
                ):
                    self.applogger.info(
                        "{}(method={}) : {} : Both access token and refresh token are expired or expiration time does not exist.".format(
                            consts.LOGS_STARTS_WITH, __method_name, self.function_name
                        )
                    )
                    self.generate_access_token(from_refresh_token)
                else:
                    self.applogger.info(
                        "{}(method={}) : {} : Using existing and valid access token.".format(
                            consts.LOGS_STARTS_WITH, __method_name, self.function_name
                        )
                    )
            else:
                self.applogger.info(
                    "{}(method={}) : {} : Access token does not exist.Generating new access token and refresh token.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
                self.generate_access_token(from_refresh_token)
        except Exception as err:
            self.applogger.error(
                "{}(method={}) : {} : Error in getting access token. : {}".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, err
                )
            )
            raise Exception(err)

    def is_token_expired(self, token_expiration):
        """To check if token is expired or not.

        Returns:
            bool: True or False
        """
        if not token_expiration:
            return True
        current_time = datetime.datetime.now().isoformat()
        return bool(token_expiration and (current_time > (token_expiration)))

    def update_checkpoint_of_disabling_function(self, make_count_zero=False):
        """Update checkpoint file for disabling function.

        Args:
            make_count_zero (bool): True when count of status code should be made zero

        Raises:
            Exception: if anything goes unintended
        """
        try:
            __method_name = inspect.currentframe().f_code.co_name
            checkpoint_file = self.state.get()
            checkpoint_json_data = {"disable_function": False}
            if checkpoint_file:
                checkpoint_json_data.update(json.loads(checkpoint_file))
            if checkpoint_json_data.get("disable_function"):
                checkpoint_json_data.update({"disable_function": False, "status_code_count": 0})
                self.state.post(json.dumps(checkpoint_json_data))
            if make_count_zero:
                if checkpoint_json_data.get("status_code_count") != 0:
                    checkpoint_json_data.update({"status_code_count": 0})
                    self.state.post(json.dumps(checkpoint_json_data))
                    self.applogger.debug(
                        "{}(method={}): {}: Checkpoint file updated as total count of status code made zero.".format(
                            consts.LOGS_STARTS_WITH, __method_name, self.function_name
                        )
                    )
                else:
                    self.applogger.debug(
                        "{}(method={}): {}: Checkpoint file already have total count of status code equal to zero.".format(
                            consts.LOGS_STARTS_WITH, __method_name, self.function_name
                        )
                    )
            else:
                current_count = checkpoint_json_data.get("status_code_count", 0)
                updated_status_code_count = current_count + 1
                if updated_status_code_count >= 10:
                    self.applogger.info(
                        "{}(method={}): {}: Function App will be Disabled.".format(
                            consts.LOGS_STARTS_WITH, __method_name, self.function_name
                        )
                    )
                    checkpoint_json_data.update({"disable_function": True, "status_code_count": updated_status_code_count})
                    self.state.post(json.dumps(checkpoint_json_data))
                    self.disable_function()
                    return

                checkpoint_json_data.update({"status_code_count": updated_status_code_count})
                self.state.post(json.dumps(checkpoint_json_data))
                self.applogger.debug(
                    "{}(method={}): {}: Updated failed status code count in checkpoint file.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )

        except Exception as err:
            self.applogger.error(
                "{}(method={}): {}: Error Occured while updating checkpoint file of disabling function. : '{}'".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, err
                )
            )

            raise Exception(f"Error Occured while updating checkpoint file of disabling function: {err}")

    def disable_function(self):
        """Disable a function app as failure count reached to 10.

        Raises:
            Exception: if any error occurs
        """
        try:
            __method_name = inspect.currentframe().f_code.co_name
            url = (consts.DISABLE_FUNCTION_APP_URL).format(consts.SUBSCRIPTION_ID, consts.RESOURCE_GROUP, consts.FUNCTION_APP_NAME)

            access_token = self.azure_token_generator()
            headers = {"Authorization": f"Bearer {access_token}"}
            response = requests.post(url=url, headers=headers)

            response.raise_for_status()
            if response.status_code in [200, 201]:
                self.applogger.info(
                    "{}(method={}): {}: Function App disabled successfully.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
            else:
                self.applogger.error(
                    "{}(method={}): {}: Error occurred while disabling function. status_code={}, reason={}".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name, response.status_code, response.text
                    )
                )
                raise Exception
        except requests.exceptions.HTTPError as err:
            self.applogger.error(
                "{}(method={}): {}: Error occurred while disabling function : status code: {}, reason: '{}'.".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, response.status_code, err
                )
            )
            self.save_checkpoint(data={"disable_function": False}, disable_function_error=True)
            raise Exception
        except Exception as err:
            self.applogger.error(
                "{}(method={}): {}: Error occured while disabling function : '{}'.".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, err
                )
            )
            self.save_checkpoint(data={"disable_function": False}, disable_function_error=True)
            raise Exception

    def azure_token_generator(self):
        """Generate access token for azure."""
        try:
            __method_name = inspect.currentframe().f_code.co_name
            url = consts.AZURE_AUTHENTICATION_URL.format(consts.AZURE_TENANT_ID)

            data = {
                "client_id": consts.AZURE_CLIENT_ID,
                "client_secret": consts.AZURE_CLIENT_SECRET,
                "grant_type": "client_credentials",
                "scope": consts.AZURE_AUTHENTICATION_SCOPE,
            }

            response = requests.get(url=url, data=data)
            if response.status_code in [200, 201]:
                self.applogger.info(
                    "{}(method={}): {}: Azure Token Generated Successfully.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
                return response.json()["access_token"]
            else:
                self.applogger.error(
                    "{}(method={}): {}: Unknown Status code while generating access token for AZURE: {} .".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name, response.status_code
                    )
                )
                self.save_checkpoint(data={"disable_function": False}, disable_function_error=True)
                raise Exception

        except Exception as err:
            self.applogger.info(
                "{}(method={}): {}: Error occured while generating access token for AZURE: {}.".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, err
                )
            )
            self.save_checkpoint(data={"disable_function": False}, disable_function_error=True)
            raise Exception

    def validate_params(self, client_id_name, client_secret_name, snapshot=False):
        """To validate parameters of function app."""
        __method_name = inspect.currentframe().f_code.co_name
        required_params = {
            "BaseURL": self.base_url,
            client_id_name: self.client_id,
            client_secret_name: self.client_secret,
            "WorkspaceID": consts.WORKSPACE_ID,
            "WorkspaceKey": consts.WORKSPACE_KEY,
            "Detections_Table_Name": consts.DETECTIONS_TABLE_NAME,
            "Audits_Table_Name": consts.AUDITS_TABLE_NAME,
            "Entity_Scoring_Table_Name": consts.ENTITY_SCORING_TABLE_NAME,
            "Lockdown_Table_Name": consts.LOCKDOWN_TABLE_NAME,
            "Health_Table_Name": consts.HEALTH_TABLE_NAME,
            "AZURE_CLIENT_ID": consts.AZURE_CLIENT_ID,
            "AZURE_CLIENT_SECRET": consts.AZURE_CLIENT_SECRET,
            "AZURE_TENANT_ID": consts.AZURE_TENANT_ID,
            "KeyVaultName": consts.KEYVAULT_NAME,
            "Entities_Table_Name": consts.ENTITIES_TABLE_NAME,
            "Azure_Subscription_Id": consts.SUBSCRIPTION_ID,
            "Azure_Resource_Group_Name": consts.RESOURCE_GROUP,
        }
        missing_required_field = False
        for label, params in required_params.items():
            if not params:
                missing_required_field = True
                self.applogger.error(
                    '{}(method={}): {}: "{}" field is not configured. field_value="{}"'.format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        label,
                        params,
                    )
                )
        if missing_required_field:
            raise Exception("Error Occurred while validating params. Required fields missing.")
        if not self.base_url.startswith("https://"):
            self.applogger.error(
                '{}(method={}) : {} : "BaseURL" must start with ”https://” schema followed '
                'by hostname. BaseURL="{}"'.format(
                    consts.LOGS_STARTS_WITH,
                    __method_name,
                    self.function_name,
                    self.base_url,
                )
            )
            raise Exception("Error Occurred while validating params. Invalid format for BaseURL.")

        if not snapshot:
            try:
                if self.start_time:
                    input_date = datetime.datetime.strptime(self.start_time, r"%m/%d/%Y %H:%M:%S")
                    now = datetime.datetime.utcnow()
                    if input_date > now:
                        self.applogger.error(
                            '{}(method={}) : {} : "StartTime" should not be in future. StartTime="{}"'.format(
                                consts.LOGS_STARTS_WITH,
                                __method_name,
                                self.function_name,
                                self.start_time,
                            )
                        )
                        raise Exception("Error Occurred while validating params. StartTime cannot be in the future.")
                    self.start_time = datetime.datetime.strftime(input_date, r"%Y-%m-%dT%H:%M:%SZ")
                else:
                    self.start_time = (
                        datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc) - datetime.timedelta(hours=24)
                    ).strftime(r"%Y-%m-%dT%H:%M:%SZ")
            except ValueError:
                self.applogger.error(
                    '{}(method={}) : {} : "StartTime" should be in "MM/DD/YYYY HH:MM:SS" (UTC) '
                    'format. StartTime="{}"'.format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        self.start_time,
                    )
                )
                raise Exception("Error Occurred while validating params. Invalid StartTime format.")

    def validate_connection(self):
        """To validate the connection with vectra and generate access token."""
        __method_name = inspect.currentframe().f_code.co_name
        try:
            self.load_tokens_from_keyvault()
            self.load_expiration_time_from_checkpoint_file()
            self.validate_access_token()
            if self.access_token:
                self.session.headers["Authorization"] = "bearer {}".format(self.access_token)
                self.applogger.info(
                    "{}(method={}) : {} : Token or Connection validation is successful.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
                return
        except Exception as ex:
            raise Exception(ex)

    def pull_data(self, endpoint, params=None):
        """To call vectra API till retry condition is not met.

        Args:
            endpoint (str): endpoint to call in Vectra.
            params (dict, optional): params to pass to API. Defaults to None.

        Returns:
            response: response object
        """
        retry_flag = True
        retry_count = 0
        while retry_flag and retry_count <= 1:
            retry_flag, res = self.pull(url=self.base_url + endpoint, params=params)
            retry_count += 1
        if retry_count > 1 and not res:
            self.update_checkpoint_of_disabling_function()
            raise Exception("Retry limit exceeded. Hence exiting the function.")
        return res

    def pull(self, url, params=None, data=None, auth=None, method="GET"):
        """Will pull data from each API.

        Args:
            url (str): _description_
            params (json, optional): _description_. Defaults to None.
            data (json, optional): _description_. Defaults to None.
            auth (json, optional): _description_. Defaults to None.
            method (str, optional): _description_. Defaults to "GET".

        Raises:
            Exception: _description_
            Exception: _description_
            Exception: _description_

        Returns:
            _type_: _description_
        """
        try:
            __method_name = inspect.currentframe().f_code.co_name
            self.applogger.debug(
                "{}(method={}) : {} request: url='{}' version='{}'".format(
                    consts.LOGS_STARTS_WITH,
                    __method_name,
                    self.function_name,
                    url,
                    self.session.headers["User-Agent"],
                )
            )
            if method == "POST":
                res = self.session.post(
                    url=url,
                    auth=auth,
                    params=params,
                    data=data,
                    timeout=consts.API_TIMEOUT,
                )
            else:
                res = self.session.get(
                    url=url,
                    auth=auth,
                    params=params,
                    data=data,
                    timeout=consts.API_TIMEOUT,
                )

            if res.status_code == 401:
                self.applogger.warning(
                    "{}(method={}) : {} : Access token is invalid. generating new token.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
                if not self.is_token_expired(self.access_token_expiration) and not self.is_token_expired(
                    self.refresh_token_expiration
                ):
                    self.generate_access_token(from_refresh_token=False)
                else:
                    self.validate_access_token()
                return True, None
            elif res and res.status_code in [200, 201]:
                self.applogger.debug(
                    '{}(method={}) : {} : API call: Response received successfully. url="{}" params="{}"'.format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        url,
                        params,
                    )
                )
                self.update_checkpoint_of_disabling_function(make_count_zero=True)
                return False, res.json()

            elif res.status_code in [400, 403, 404]:
                self.applogger.error(
                    "{}(method={}) : {} : API call: response: url={} status_code={} response={}".format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        url,
                        res.status_code,
                        res.text,
                    )
                )
                self.update_checkpoint_of_disabling_function()
                raise Exception
            else:
                self.applogger.error(
                    "{}(method={}) : {} : API call: Unknown status code or empty "
                    'response: url="{}" status_code="{}" response="{}"'.format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        url,
                        res.status_code,
                        res.text,
                    )
                )
                raise Exception("Received unknown status code or empty response.")
        except requests.exceptions.RequestException as ex:
            if res.status_code == 404:
                self.applogger.debug(
                    '{}(method={}) : {} : API call: Got {} Status Code : url="{}"'
                    ' response="{}"'.format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        url,
                        ex.response.status_code,
                        ex.response.text,
                    )
                )
                return {}
            else:
                self.applogger.error(
                    '{}(method={}) : {} : API call: Unsuccessful response: url="{}" status_code="{}"'
                    ' response="{}"'.format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        url,
                        ex.response.status_code,
                        ex.response.text,
                    )
                )
                raise Exception("HTTP Error Occurred while getting response from api.")
        except Exception as ex:
            self.applogger.error(
                '{}(method={}) : {} : API call: Unexpected error while API call url="{}" error="{}"'.format(
                    consts.LOGS_STARTS_WITH,
                    __method_name,
                    self.function_name,
                    url,
                    str(ex),
                )
            )
            raise Exception("Error Occurred while getting response from api.")

    def get_checkpoint_field_and_value(self):
        """Fetch last data from checkpoint file.

        Returns:
            None/json: last_data
        """
        try:
            __method_name = inspect.currentframe().f_code.co_name
            field = None
            checkpoint = self.state.get()
            if checkpoint:
                checkpoint = json.loads(checkpoint)
                if checkpoint.get("from", None):
                    field, checkpoint = "from", checkpoint.get("from")
                elif checkpoint.get("event_timestamp_gte"):
                    field, checkpoint = "event_timestamp_gte", checkpoint.get("event_timestamp_gte")
                else:
                    field, checkpoint = "event_timestamp_gte", self.start_time
                    self.state.post(json.dumps({"event_timestamp_gte": self.start_time}))
            else:
                field, checkpoint = "event_timestamp_gte", self.start_time
                self.state.post(json.dumps({"event_timestamp_gte": self.start_time}))
            self.applogger.info(
                '{}(method={}) : {} : Checkpoint field="{}" and value="{}"'.format(
                    consts.LOGS_STARTS_WITH,
                    __method_name,
                    self.function_name,
                    field,
                    checkpoint,
                )
            )
            return field, checkpoint
        except Exception as ex:
            self.applogger.error(
                '{}(method={}) : {} : Unexpected error while getting checkpoint: err="{}"'.format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, str(ex)
                )
            )
            raise Exception(ex)

    def save_checkpoint(
        self, data, token_expiry_time=False, entity_state=None, entities_endpoint=False, disable_function_error=False
    ):
        """Post checkpoint into sentinel."""
        try:
            __method_name = inspect.currentframe().f_code.co_name
            checkpoint_data = self.state.get()
            checkpoint_data_json = {}
            if checkpoint_data:
                checkpoint_data_json = json.loads(checkpoint_data)
            if token_expiry_time:
                checkpoint_data_json.update(data)
                self.state.post(json.dumps(checkpoint_data_json))
                self.applogger.info(
                    '{}(method={}) : {} : successfully saved updated expiration time in checkpoint."{}"'.format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        data,
                    )
                )
            elif entities_endpoint:
                checkpoint_data_json.update({"last_modified_timestamp": data})
                entity_state.post(json.dumps(checkpoint_data_json))
                self.applogger.info(
                    '{}(method={}) : {} : successfully saved checkpoint. last_modified_timestamp="{}"'.format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        data,
                    )
                )
            elif disable_function_error:
                checkpoint_data_json.update(data)
                self.state.post(json.dumps(checkpoint_data_json))
                self.applogger.debug(
                    '{}(method={}) : {} : error while disabling function. So making disable_function flag to False.'.format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name
                    )
                )
            else:
                checkpoint_data_json.update({"from": data})
                self.state.post(json.dumps(checkpoint_data_json))
                self.applogger.info(
                    '{}(method={}) : {} : successfully saved checkpoint. from="{}"'.format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        data,
                    )
                )
        except Exception as ex:
            self.applogger.exception(
                '{}(method={}) : {} : Unexpected error while saving checkpoint: err="{}"'.format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, str(ex)
                )
            )
            raise Exception(ex)

    def post_data_to_sentinel(self, data, table_name, fields):
        """To post data into sentinel."""
        __method_name = inspect.currentframe().f_code.co_name
        if fields:
            for event in data:
                for field in fields:
                    event[field] = [event.get(field)]
        data = json.dumps(data)
        status_code = self.azuresentinel.post_data(data, table_name)
        if status_code in consts.SENTINEL_ACCEPTABLE_CODES:
            self.applogger.info(
                '{}(method={}) : {} : Successfully posted the data in the table="{}"'.format(
                    consts.LOGS_STARTS_WITH,
                    __method_name,
                    self.function_name,
                    table_name,
                )
            )
        else:
            self.applogger.error(
                '{}(method={}) : {} : Data is not posted in the table="{}" status_code="{}"'.format(
                    consts.LOGS_STARTS_WITH,
                    __method_name,
                    self.function_name,
                    table_name,
                    status_code,
                )
            )
            raise Exception("Error Occurred while posting data into Microsoft Sentinel Log Analytics Workspace.")

    def _get_size_of_chunk_in_mb(self, chunk):
        """Get the size of chunk in MB."""
        return sys.getsizeof(chunk) / (1024 * 1024)

    def _create_chunks_and_post_to_sentinel(
        self, data, table_name, fields, entity_state=None, entities_endpoint=False
    ):
        """Create chunks and post to chunk to sentinel."""
        __method_name = inspect.currentframe().f_code.co_name
        chunk = []
        if self._get_size_of_chunk_in_mb(data) < 30:
            self.post_data_to_sentinel(data, table_name, fields)
            return
        for event in data:
            chunk.append(event)
            if self._get_size_of_chunk_in_mb(chunk) >= 30:
                if chunk[:-1]:
                    self.post_data_to_sentinel(chunk[:-1], table_name, fields)
                    if entities_endpoint:
                        next_checkpoint = chunk[-2].get("last_modified_timestamp")

                    else:
                        next_checkpoint = chunk[-2].get("id")
                    self.save_checkpoint(
                        data=next_checkpoint, entity_state=entity_state, entities_endpoint=entities_endpoint
                    )
                    chunk = [event]
                    continue
                else:
                    id = chunk[0].get("id")
                    self.applogger.error(
                        "{}(method={}) : {} : event with id {} is too large to post into the sentinel hence skipping it.".format(
                            consts.LOGS_STARTS_WITH,
                            __method_name,
                            self.function_name,
                            id,
                        )
                    )
                    chunk = []
                    continue
        if chunk:
            self.post_data_to_sentinel(chunk, table_name, fields)

    def pull_and_push_the_data(
        self,
        endpoint,
        checkpoint_field,
        checkpoint_value,
        table_name,
        fields=None,
        query_params=dict(),
    ):
        """To pull the data from vectra and push into sentinel."""
        __method_name = inspect.currentframe().f_code.co_name
        posted_event_count = 0
        iter_next = True
        query_params.update({"limit": consts.PAGE_SIZE, checkpoint_field: checkpoint_value})
        while iter_next:
            res = self.pull_data(endpoint, params=query_params)
            next_checkpoint = res.get("next_checkpoint", None)
            if endpoint == consts.DETECTIONS_ENDPOINT and len(res.get("events")):
                self.applogger.debug(
                    "{}(method={}) : {} : Trying to collect the additional information from"
                    " /detections endpoint for type=host and type=account.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
                detection_ids_list = []
                for event in res.get("events"):
                    if event.get("type") == "host" or event.get("type") == "account":
                        detection_ids_list.append(str(event.get("detection_id")))

                detection_id_set = set(detection_ids_list)
                detections_ids = ",".join(detection_id_set)
                next = True
                page = 1
                merge_json = {}

                if detections_ids:
                    while next:
                        detection_endpoint = "/api/v3.3/detections"
                        self.applogger.debug(
                            "{}(method={}) : {} : GET call to /detections endpoint with URL = {}".format(
                                consts.LOGS_STARTS_WITH,
                                __method_name,
                                self.function_name,
                                self.base_url + detection_endpoint,
                            )
                        )
                        host_details = self.pull_data(detection_endpoint, params={"id": detections_ids, "page": page})
                        if host_details.get("results"):
                            for each in host_details.get("results"):
                                merge_json[each.get("id")] = each
                        if not host_details.get("next"):
                            break
                        page += 1

                for event in res.get("events"):
                    if event.get("type") == "host" and merge_json.get(event.get("detection_id"), {}):
                        temp_host = merge_json.get(event.get("detection_id"), {})
                        event["d_detection_details"] = [temp_host]
                        event["is_targeting_key_asset"] = str(temp_host.get("is_targeting_key_asset", ""))
                        event["src_host"] = [temp_host.get("src_host", {})]
                        event["normal_domains"] = temp_host.get("normal_domains", [])
                        event["src_ip"] = temp_host.get("src_ip", "")
                        event["summary"] = [temp_host.get("summary", {})]
                        event["grouped_details"] = temp_host.get("grouped_details", [])
                        event["tags"] = temp_host.get("tags", [])
                    elif event.get("type") == "account":
                        if event.get("detail", {}):
                            event["d_detection_details"] = [event.get("detail", {})]
                        if merge_json.get(event.get("detection_id"), {}):
                            account_data = merge_json.get(event.get("detection_id"), {})
                            event["tags"] = account_data.get("tags", [])
                    else:
                        event["d_detection_details"] = []
                    self.applogger.debug(
                        "{}(method={}) : {} :  Successfully modified events/detections"
                        " response for id={}.".format(
                            consts.LOGS_STARTS_WITH,
                            __method_name,
                            self.function_name,
                            event.get("id"),
                        )
                    )

            if res and len(res.get("events")):
                self._create_chunks_and_post_to_sentinel(res.get("events"), table_name, fields)
                posted_event_count += len(res.get("events"))
                iter_next = True if int(res.get("remaining_count")) > 0 else False
                query_params.update({"limit": consts.PAGE_SIZE, "from": next_checkpoint})
            else:
                iter_next = False
            if endpoint == consts.ENTITY_SCORING_ENDPOINT and (next_checkpoint is None or next_checkpoint == "null"):
                break
            else:
                self.save_checkpoint(data=next_checkpoint)

        self.applogger.info(
            "{}(method={}) : {} : Posted total {} event(s) into MS Sentinel. No more events."
            " Stopping the collection.".format(
                consts.LOGS_STARTS_WITH,
                __method_name,
                self.function_name,
                posted_event_count,
            )
        )

    def pull_and_push_the_snapshot_data(
        self,
        endpoint,
        table_name,
        hashed_events_list=list(),
        hash_field_list=[],
        fields=None,
    ):
        """To pull the snapshot type data from vectra and push into sentinel."""
        __method_name = inspect.currentframe().f_code.co_name
        posted_event_count = 0
        res = self.pull_data(endpoint)
        if res and len(res):
            if endpoint == consts.HEALTH_ENDPOINT:
                link_status_dict, aggregated_peak_traffic_dict = {}, {}
                connectivity_dict, trafficdrop_dict = {}, {}

                # for link_status
                for k, v in res.get("network", {}).get("interfaces", {}).get("sensors", {}).items():
                    link_status = "UP"
                    for x, y in v.items():
                        if y.get("link", "") != "UP":
                            link_status = "Degraded"
                            break
                    link_status_dict[k] = link_status

                # for aggregated_peak_traffic
                for key, value in res.get("network", {}).get("traffic", {}).get("sensors", {}).items():
                    aggregated_peak_traffic_dict[key] = value.get("aggregated_peak_traffic_mbps", "")

                # for connectivity status and error
                for item in res.get("connectivity", {}).get("sensors", {}):
                    connectivity_dict[item.get("name", "")] = {
                        "status": item.get("status", ""),
                        "error": item.get("error", ""),
                    }

                # for traffic_drop status and error
                for item in res.get("trafficdrop", {}).get("sensors", {}):
                    trafficdrop_dict[item.get("name", "")] = {
                        "status": item.get("status", ""),
                        "error": item.get("error", ""),
                    }

                for i in range(len(res.get("sensors", {}))):
                    # adding d_link_status
                    res["sensors"][i]["d_link_status"] = link_status_dict.get(
                        res.get("sensors", {})[i].get("luid", ""), ""
                    )
                    # adding d_aggregated_peak_traffic
                    res["sensors"][i]["d_aggregated_peak_traffic"] = aggregated_peak_traffic_dict.get(
                        res.get("sensors", {})[i].get("name", ""), ""
                    )
                    # adding d_connectivity_status
                    res["sensors"][i]["d_connectivity_status"] = connectivity_dict.get(
                        res.get("sensors", {})[i].get("name", ""), {}
                    ).get("status", "")
                    # adding d_connectivity_error
                    res["sensors"][i]["d_connectivity_error"] = connectivity_dict.get(
                        res.get("sensors", {})[i].get("name", ""), {}
                    ).get("error", "")
                    # adding d_trafficdrop_status
                    res["sensors"][i]["d_trafficdrop_status"] = trafficdrop_dict.get(
                        res.get("sensors", {})[i].get("name", ""), {}
                    ).get("status", "")
                    # adding d_trafficdrop_error
                    res["sensors"][i]["d_trafficdrop_error"] = trafficdrop_dict.get(
                        res.get("sensors", {})[i].get("name", ""), {}
                    ).get("error", "")

                list_res = [res]
                self.post_data_to_sentinel(list_res, table_name, fields)
                posted_event_count += 1
            else:
                for event in res:
                    temp_dict = {}
                    for field in hash_field_list:
                        temp_dict[field] = event.get(field)
                    hash_of_event = self.get_results_hash(temp_dict)
                    if hash_of_event not in hashed_events_list:
                        self.post_data_to_sentinel(event, table_name, fields)
                        posted_event_count += 1
                        hashed_events_list.append(hash_of_event)
                        self.save_checkpoint_snapshot(hashed_events_list)

        self.applogger.info(
            "{}(method={}) : {} : Posted total {} event(s) into MS Sentinel. No more events."
            " Stopping the collection.".format(
                consts.LOGS_STARTS_WITH,
                __method_name,
                self.function_name,
                posted_event_count,
            )
        )

    def get_results_hash(self, data):
        """Generate hash digest.

        :data: Data to be hashed.
        :return: SHA512 hexdigest of data.
        """
        data = json.dumps(data, sort_keys=True)
        result = hashlib.sha512(data.encode())
        result_hash = result.hexdigest()
        return result_hash

    def get_checkpoint_snapshot(self):
        """Fetch snapshot hash from checkpoint file.

        Returns:
            List: hash_list
        """
        try:
            __method_name = inspect.currentframe().f_code.co_name
            checkpoint = self.state.get()
            if checkpoint:
                checkpoint = json.loads(checkpoint)
                checkpoint = checkpoint.get("snapshot")
                self.applogger.info(
                    "{}(method={}) : {} : Checkpoint list fetched successfully. checkpoint={}".format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        checkpoint,
                    )
                )
            else:
                checkpoint = []
                self.state.post(json.dumps({"snapshot": checkpoint}))
                self.applogger.info(
                    "{}(method={}) : {} : Checkpoint list not found. Created new checkpoint list.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
            return checkpoint
        except Exception as ex:
            self.applogger.error(
                '{}(method={}) : {} : Unexpected error while getting checkpoint list: err="{}"'.format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, str(ex)
                )
            )
            raise Exception(ex)

    def save_checkpoint_snapshot(self, value):
        """Post checkpoint snapshot into sentinel."""
        try:
            __method_name = inspect.currentframe().f_code.co_name
            checkpoint_data = self.state.get()
            checkpoint_data_json = {}
            if checkpoint_data:
                checkpoint_data_json = json.loads(checkpoint_data)
            checkpoint_data_json.update({"snapshot": value})
            self.state.post(json.dumps(checkpoint_data_json))
            self.applogger.info(
                "{}(method={}) : {} : successfully saved checkpoint.".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name
                )
            )
        except Exception as ex:
            self.applogger.exception(
                '{}(method={}) : {} : Unexpected error while saving checkpoint: err="{}"'.format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, str(ex)
                )
            )
            raise Exception(ex)

    def get_entities_checkpoint(self, entity_state):
        """Fetch last data from checkpoint file for entities.

        Args:
            entity_state (object): State manager object

        Returns:
            None/json: last_data
        """
        try:
            __method_name = inspect.currentframe().f_code.co_name
            self.applogger.debug(
                "{}(method={}) : {} : Fetching last data from checkpoint file.".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name
                )
            )
            checkpoint = entity_state.get()
            if checkpoint:
                checkpoint = json.loads(checkpoint)
                if checkpoint.get("last_modified_timestamp", None):
                    checkpoint = checkpoint.get("last_modified_timestamp")
                    self.applogger.info(
                        '{}(method={}) : {} : Checkpoint value="{}"'.format(
                            consts.LOGS_STARTS_WITH,
                            __method_name,
                            self.function_name,
                            checkpoint,
                        )
                    )
                    return checkpoint
            return None
        except Exception as ex:
            self.applogger.error(
                '{}(method={}) : {} : Unexpected error while getting checkpoint: err="{}"'.format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, str(ex)
                )
            )
            raise Exception(ex)

    def merge_assignment_details(self, entity_response, entity_type):
        """
        To merge assignment details.

        Args:
            entity_response (list): List of entities
            entity_type (str): Type of the entity (host/account)

        Returns:
            list: Updated List of entities
        """
        __method_name = inspect.currentframe().f_code.co_name
        try:
            self.applogger.debug(
                "{}(method={}) : {} : Trying to collect information from {} endpoint.".format(
                    consts.LOGS_STARTS_WITH,
                    __method_name,
                    self.function_name,
                    consts.ASSIGNMENT_ENDPOINT,
                )
            )
            entity_ids_list = [str(event.get("id")) for event in entity_response]
            entity_ids = ",".join(entity_ids_list)
            entity_type_params = entity_type + "s"
            merge_json = {}
            if entity_ids:
                assignment_params = {
                    "resolved": False,
                    "page_size": consts.PAGE_SIZE,
                    entity_type_params: entity_ids,
                }
                assignment_data = self.pull_data(
                    endpoint=consts.ASSIGNMENT_ENDPOINT, params=assignment_params
                )
                self.applogger.debug(
                    "{}(method={}) : {} : Collected information from {} endpoint.".format(
                        consts.LOGS_STARTS_WITH,
                        __method_name,
                        self.function_name,
                        consts.ASSIGNMENT_ENDPOINT,
                    )
                )
                if assignment_data.get("results"):
                    for assignment_result in assignment_data.get("results"):
                        merge_json[assignment_result.get(entity_type + "_id")] = assignment_result

                for event in entity_response:
                    event['url'] = consts.BASE_URL + '/{}/{}?pivot={}'.format(entity_type_params, event.get("id"), consts.USER_AGENT)
                    if merge_json.get(event.get("id"), {}):
                        assignment_details = merge_json.get(event.get("id"), {})
                        entity_assignment = assignment_details.get("assigned_to", {})
                        event["assigned_to_username"] = entity_assignment.get(
                            "username", ""
                        )
            return entity_response
        except Exception as ex:
            self.applogger.error(
                '{}(method={}) : {} : Unexpected error while getting assignments data: err="{}"'.format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, str(ex)
                )
            )
            raise Exception(ex)

    def pull_and_push_entities_data(
        self,
        endpoint,
        checkpoint_value,
        table_name,
        entity_type,
        entity_state,
        fields=None,
    ):
        """
        To pull the entities data from vectra and push into sentinel.

        Args:
            endpoint (str): endpoint url
            checkpoint_value (str): Checkpoint value from checkpoint file
            table_name (str): Table name to ingest data into sentinel
            entity_type (str): Type of the entity (host/account)
            entity_state (object): State manager object
            fields (list): list of fields
        """
        __method_name = inspect.currentframe().f_code.co_name
        try:
            self.applogger.debug(
                "{}(method={}) : {} Pull and Push entities data".format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name
                )
            )
            posted_entities_count = 0
            page = 1
            params = {
                "page_size": consts.PAGE_SIZE,
                "ordering": "last_modified_timestamp",
                "type": entity_type,
            }

            if checkpoint_value:
                params.update({"last_modified_timestamp_gte": checkpoint_value})
            while True:
                params.update({"page": page})
                res = self.pull_data(endpoint, params=params)
                entity_response = res.get("results")
                self.applogger.debug(
                    "{}(method={}) : {} : Trying to collect information from"
                    " /entities endpoint.".format(
                        consts.LOGS_STARTS_WITH, __method_name, self.function_name
                    )
                )
                next_link = res.get("next", None)
                if len(entity_response):
                    last_modified_timestamp = entity_response[-1].get(
                        "last_modified_timestamp"
                    )
                    merged_entity_response = self.merge_assignment_details(
                        entity_response, entity_type
                    )
                    self._create_chunks_and_post_to_sentinel(
                        merged_entity_response,
                        table_name,
                        fields,
                        entity_state,
                        entities_endpoint=True,
                    )
                    posted_entities_count += len(entity_response)
                    self.applogger.info(
                        "{}(method={}) : {} : Posted Entities Count : {}".format(
                            consts.LOGS_STARTS_WITH,
                            __method_name,
                            self.function_name,
                            posted_entities_count,
                        )
                    )
                    self.save_checkpoint(
                        data=last_modified_timestamp, entity_state=entity_state, entities_endpoint=True
                    )
                if not next_link:
                    break
                page += 1
        except Exception as ex:
            self.applogger.error(
                '{}(method={}) : {} : Unexpected error while getting or posting entities data: err="{}"'.format(
                    consts.LOGS_STARTS_WITH, __method_name, self.function_name, str(ex)
                )
            )
            raise Exception(ex)